import torch
import torch.utils.data as data
import torch.nn as nn
import torch.optim as optim
import numpy as np
import pandas as pd
import cv2
import sys

# 参数初始化
def gaussian_weights_init(m):
    classname = m.__class__.__name__
    # 字符串查找find，找不到返回-1，不等-1即字符串中含有该字符
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.04)

class FaceCNN(nn.Module):
    # 初始化网络结构
    def __init__(self):
        super(FaceCNN, self).__init__()

        # 第一次卷积、池化
        self.conv1 = nn.Sequential(
            # 输入通道数in_channels，输出通道数(即卷积核的通道数)out_channels，卷积核大小kernel_size，步长stride，对称填0行列数padding
            # input:(bitch_size, 1, 48, 48), output:(bitch_size, 64, 48, 48), (48-3+2*1)/1+1 = 48
            nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, stride=1, padding=1), # 卷积层
            nn.BatchNorm2d(num_features=64), # 归一化
            nn.RReLU(inplace=True), # 激活函数
            # output(bitch_size, 64, 24, 24)
            nn.MaxPool2d(kernel_size=2, stride=2), # 最大值池化
        )

        # 第二次卷积、池化
        self.conv2 = nn.Sequential(
            # input:(bitch_size, 64, 24, 24), output:(bitch_size, 128, 24, 24), (24-3+2*1)/1+1 = 24
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(num_features=128),
            nn.RReLU(inplace=True),
            # output:(bitch_size, 128, 12 ,12)
            nn.MaxPool2d(kernel_size=2, stride=2),
        )

        # 第三次卷积、池化
        self.conv3 = nn.Sequential(
            # input:(bitch_size, 128, 12, 12), output:(bitch_size, 256, 12, 12), (12-3+2*1)/1+1 = 12
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(num_features=256),
            nn.RReLU(inplace=True),
            # output:(bitch_size, 256, 6 ,6)
            nn.MaxPool2d(kernel_size=2, stride=2),
        )

        # 参数初始化
        self.conv1.apply(gaussian_weights_init)
        self.conv2.apply(gaussian_weights_init)
        self.conv3.apply(gaussian_weights_init)

        # 全连接层
        self.fc = nn.Sequential(
            nn.Dropout(p=0.2),
            nn.Linear(in_features=256*6*6, out_features=4096),
            nn.RReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(in_features=4096, out_features=1024),
            nn.RReLU(inplace=True),
            nn.Linear(in_features=1024, out_features=256),
            nn.RReLU(inplace=True),
            nn.Linear(in_features=256, out_features=7),
        )

    # 前向传播
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        # 数据扁平化
        x = x.view(x.shape[0], -1)
        y = self.fc(x)
        return y

def webcam(path):  #'./video/114514.mp4',将老师给的代码打包成一个子函数
    #传统处理初始化
    facecascPath = "src/Webcam-Face-Detect/data/haarcascade_frontalface_default.xml"
    left_eyecascPath = "src/Webcam-Face-Detect/data/haarcascade_lefteye_2splits.xml"
    right_eyecascPath = "src/Webcam-Face-Detect/data/haarcascade_righteye_2splits.xml"
    faceCascade = cv2.CascadeClassifier(facecascPath)
    left_eyeCasecade = cv2.CascadeClassifier(left_eyecascPath)
    right_eyeCasecade = cv2.CascadeClassifier(right_eyecascPath)
    video_capture = cv2.VideoCapture(path)

    #获取视频信息
    fps = int(video_capture.get(cv2.CAP_PROP_FPS))
    frame_width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
     # 设置视频写入对象
    video_out = cv2.VideoWriter('/home/misaka/opencvshow/video/output/output_video.avi', cv2.VideoWriter_fourcc(*'XVID'), fps, (frame_width, frame_height))

    #模型初始化
    #使用自带的rss模型
    detection_model_path = '/home/misaka/opencvshow/src/Webcam-Face-Detect/data/haarcascade_frontalcatface_extended.xml'
    classification_model_path = '/home/misaka/opencvshow/src/Webcam-Face-Detect/data/dataset/model/model_cnn.pkl'
    #传统的人脸检测模型
    face_detection = cv2.CascadeClassifier(detection_model_path)
    #加载模型
    emotion_classifer = torch.load(classification_model_path)
    
    emotion_labels = ['Where is your mother?!', 'Shift', 'Where is my mom?', 'Happy happy happy', 'Ing ing ing', 'What the fuck', 'Calm as dead']
    #emotion_labels = ['angry', 'heat', 'fear', 'Happy', 'Sad', 'puzzle', 'Calm']
    

    while video_capture.isOpened():
        # Capture frame-by-frame
        ret, frame = video_capture.read()


        if not ret:
            print("End of video")
            break




        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        faces = faceCascade.detectMultiScale(
            gray,
            scaleFactor=1.1,
            minNeighbors=5,
            minSize=(30, 30),

        )

        # Draw a rectangle around the faces
        for (x, y, w, h) in faces:
            cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
            faceimage=gray[y:y+h, x:x+w]
            faceimage_roi=cv2.resize(faceimage, (48, 48))
            face_tensor = torch.tensor(faceimage_roi).unsqueeze(0).unsqueeze(0).float()
            
            outputs = emotion_classifer(face_tensor)

            probabilities = torch.softmax(outputs, dim=1)
            # 获取预测的情绪类别索引
            predicted_idx = torch.argmax(probabilities, dim=1)
            # 获取预测的情绪标签
            predicted_label = emotion_labels[predicted_idx.item()]
            #print(f"Predicted emotion: {predicted_label}")

            #显示在图像上
            frame_new=cv2.putText(frame, predicted_label, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)
            video_out.write(frame_new)
            
        cv2.imshow('Video', frame)
        cv2.imshow('face', faceimage)
        




        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    video_capture.release()
    video_out.release()
    cv2.destroyAllWindows()




def main():


    webcam('/home/misaka/opencvshow/video/114514.mp4')
    webcam('/home/misaka/opencvshow/video/text_video.mp4')
    
    #webcam('/home/misaka/opencvshow/video/1.mp4')
    webcam('/home/misaka/opencvshow/video/text_video2.mp4')
    #webcam('/home/misaka/opencvshow/video/carton.mp4')




if __name__ == '__main__':
    main()
